#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: 邵奈一
@Email: shaonaiyi@163.com
@Date: 2024/11/14
@微信：shaonaiyi888
@微信公众号: 邵奈一 
"""
# 代码4-1
# 定义生成图像集路径txt文档的函数
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
from PIL import Image


# 定义创建关于图像路径txt文件的函数
def classes_txt(root, out_path, num_class=None):
    # 列出根目录下所有类别所在文件夹名
    dirs = os.listdir(root)
    # 不指定类别数量就读取所有
    if not num_class:
        num_class = len(dirs)
    # 输出文件路径不存在就新建
    if not os.path.exists(out_path):
        f = open(out_path, 'w')
        f.close()
    # 如果文件中本来就有一部分内容，只需要补充剩余部分
    # 如果文件中数据的类别数比需要的多就跳过
    # 打开文件，以读写模式
    with open(out_path, 'r+') as f:
        # 尝试读取文件最后一行的倒数第二个字段，并将其转换为整数，然后加1
        try:
            # f.readlines()[-1]表示选取文件中的最后一行，使用 os.path.sep 来获取当前系统的路径分隔符
            end = int(f.readlines()[-1].split(os.path.sep)[-2]) + 1
        # 如果出现异常（例如文件为空，或者最后一行没有倒数第二个字段），则将end设为0
        except:
            end = 0
        # 如果end小于num_class - 1，说明需要添加新的目录和文件
        if end < num_class - 1:
            # 对dirs进行排序
            dirs.sort()
            # 从end开始，取出num_class - end个元素，作为新的目录列表
            dirs = dirs[end:num_class]
            # 遍历新的目录列表
            for dir in dirs:
                # 获取当前目录下的所有文件
                files = os.listdir(os.path.join(root, dir))
                # 遍历文件列表
                for file in files:
                    # 将文件路径写入文件
                    f.write(os.path.join(root, dir, file) + '\n')


# 代码4-2
# 定义读取并变换图像数据格式的类
# 定义读取并变换数据格式的类
class MyDataset(Dataset):
    def __init__(self, txt_path, num_class, transforms=None):
        super().__init__()
        # 存储图像的路径
        images = []
        # 图像的类别名，在本例中是数字
        labels = []
        # 打开上一步生成的txt文件
        with open(txt_path, 'r') as f:
            for line in f:
                # 只读取前 num_class个类
                # 如果当前行的倒数第二个'\\'分割后的字符串转为整数大于等于num_class，则跳出循环，使用 os.path.sep 来获取当前系统的路径分隔符
                # 比如xxxx/data4/test/00000/72.png，则取00000，判断是否大于等于num_class
                if int(line.split(os.path.sep)[-2]) >= num_class:
                    break
                # 去除行尾的换行符
                line = line.strip('\n')
                # 将处理后的行添加到images列表中
                images.append(line)
                # 将行倒数第二个'\\'分割后的字符串转为整数添加到labels列表中，使用 os.path.sep 来获取当前系统的路径分隔符
                labels.append(int(line.split(os.path.sep)[-2]))
        self.images = images
        self.labels = labels
        # 图片需要进行的格式变换，比如ToTensor()等等
        self.transforms = transforms

    def __getitem__(self, index):
        # 用PIL.Image读取图像
        image = Image.open(self.images[index]).convert('RGB')
        label = self.labels[index]
        if self.transforms is not None:
            # 进行格式变换
            image = self.transforms(image)
        return image, label

    def __len__(self):
        return len(self.labels)


# 代码4-3
# 加载图像数据
# 首先将训练集和测试集文件途径和文件名以txt保存在一个文件夹中，路径自行定义
root = 'data4'  # 文件的储存位置
classes_txt(root + '/train', root + '/train.txt')
classes_txt(root + '/test', root + '/test.txt')

# 由于数据集图片尺寸不一，因此要进行resize（重设大小）
# 将图片大小重设为 64 * 64
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([transforms.Resize((64, 64)),
                                transforms.Grayscale(),
                                transforms.ToTensor()])

# 提取训练集和测试集图片的路径生成txt文件
# num_class 选取100种汉字  提出图片和标签
train_set = MyDataset(root + '/train.txt',
                      num_class=100,
                      transforms=transform)
test_set = MyDataset(root + '/test.txt',
                     num_class=100,
                     transforms=transform)
# 放入迭代器中
train_loader = DataLoader(train_set, batch_size=50, shuffle=True)
test_loader = DataLoader(test_set, batch_size=5473, shuffle=True)
# 这里的5473 是因为测试集为5973张图片，当进行迭代时取第二批500个图片进行测试
for step, (x, y) in enumerate(test_loader):
    test_x, labels_test = x.to(device), y.to(device)